import os
import requests
from langchain.agents import initialize_agent, Tool, AgentType
from langchain.llms import DashScope
from langchain.chains.conversation.memory import ConversationBufferMemory

# 1. 工具1：数学计算（仅支持基础算术运算）
def calculate_math(expression: str) -> str:
    """
    处理数学计算，支持加减乘除和括号运算，例如"3+5*2"、"(10-3)/2"
    """
    try:
        # 仅允许安全的数学运算字符
        allowed_chars = set("0123456789+-*/(). ")
        if not all(c in allowed_chars for c in expression):
            return "错误：仅支持数字和加减乘除运算"
        
        result = eval(expression)  # 注意：生产环境需用更安全的计算库
        return f"{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"

# 2. 工具2：天气查询（使用和风天气API，需自行申请key）
def get_weather(city: str) -> str:
    """
    查询指定城市的实时天气，输入格式为城市名，例如"北京"、"上海"
    """
    weather_api_key = os.getenv("WEATHER_API_KEY")
    if not weather_api_key:
        return "错误：请设置WEATHER_API_KEY环境变量（和风天气API密钥）"
    
    try:
        # 和风天气API接口（免费版）
        url = f"https://devapi.qweather.com/v7/weather/now?location={city}&key={weather_api_key}"
        response = requests.get(url, timeout=5)
        data = response.json()
        
        if data["code"] != "200":
            return f"天气查询失败：{data.get('msg', '未知错误')}"
        
        now = data["now"]
        return f"{city}当前天气：{now['text']}，温度{now['temp']}℃，风力{now['windDir']}{now['windScale']}级"
    except Exception as e:
        return f"天气查询出错：{str(e)}"

# 3. 初始化工具列表（后续扩展只需添加新Tool）
tools = [
    Tool(
        name="Calculator",
        func=calculate_math,
        description="用于处理数学计算问题，输入应为纯数学表达式（如3+5、(10-2)*3）"
    ),
    Tool(
        name="WeatherQuery",
        func=get_weather,
        description="用于查询城市实时天气，输入应为城市名称（如北京、上海）"
    )
]

# 4. 初始化千问大模型
llm = DashScope(
    model="qwen-turbo",
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"),
    temperature=0  # 降低随机性，确保工具调用稳定
)

# 5. 初始化对话记忆（支持多轮对话上下文）
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 6. 初始化智能体（使用工具调用型Agent）
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,  # 支持多轮对话的工具调用Agent
    memory=memory,
    verbose=True,  # 输出思考过程（调试用）
    handle_parsing_errors="请检查输入格式，我无法理解你的问题"
)

# 7. 运行交互循环
def main():
    print("多工具智能体已启动（支持数学计算和天气查询，输入'退出'结束）")
    while True:
        try:
            user_input = input("你: ")
            if user_input.lower() == "退出":
                print("智能体: 再见！")
                break
            response = agent.run(user_input)
            print(f"智能体: {response}")
        except KeyboardInterrupt:
            print("\n智能体: 已退出")
            break

if __name__ == "__main__":
    main()
